/* *************************************** */
/* * Smooth - MPI Edition                * */
/* * Authors:                            * */
/* *  André Bannwart Perina      7152757 * */
/* *  Rodrigo de Freitas Pereira 7573472 * */
/* *  Thaís Emanuele dos Santos  6453087 * */
/* *************************************** */

#include "../include/mpi.h"
#include "../include/ppm_funcs.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
 
int main(int argc,char *argv[]) {
	int rank, size, tag = 1;
    MPI_Status stat;
	struct timeval tEnd, tInit;

	/* Start MPI */
    MPI_Init(&argc, &argv);

	/* Get rank and size */
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

	gettimeofday(&tInit, NULL);

	/* Root Node */
    if(!rank) {
		unsigned int width, length, segSize, lastSegSize, sCount, i;
		byte *chunk;
		ppm_seg_t *r, *g, *b, *rRef, *gRef, *bRef;

		/* Create segments */
		ppm_create(&r);
		ppm_create(&g);
		ppm_create(&b);
		ppm_create_ref(&rRef);
		ppm_create_ref(&gRef);
		ppm_create_ref(&bRef);

		/* Read file */
		ppm_get_channel("in.ppm", PPM_CHANNEL_R, PPM_FORMAT_P3, r);
		ppm_get_channel("in.ppm", PPM_CHANNEL_G, PPM_FORMAT_P3, g);
		ppm_get_channel("in.ppm", PPM_CHANNEL_B, PPM_FORMAT_P3, b);

		/* Set sizes */
		width = r->n;
		length = r->m;

		/* Calculate workshare for each node */
		segSize = (unsigned int) floor((double) length / (double) (size-1));
		lastSegSize = segSize + (length % (size-1));

		/* Add borders */
		ppm_add_border(r, PPM_BORDER_TOP, 2);
		ppm_add_border(r, PPM_BORDER_BOTTOM, 2);
		ppm_add_border(r, PPM_BORDER_LEFT, 2);
		ppm_add_border(r, PPM_BORDER_RIGHT, 2);
		ppm_add_border(g, PPM_BORDER_TOP, 2);
		ppm_add_border(g, PPM_BORDER_BOTTOM, 2);
		ppm_add_border(g, PPM_BORDER_LEFT, 2);
		ppm_add_border(g, PPM_BORDER_RIGHT, 2);
		ppm_add_border(b, PPM_BORDER_TOP, 2);
		ppm_add_border(b, PPM_BORDER_BOTTOM, 2);
		ppm_add_border(b, PPM_BORDER_LEFT, 2);
		ppm_add_border(b, PPM_BORDER_RIGHT, 2);

		/* Refresh sizes */
		width += 4;
		length += 4;

		/* Broadcast image width */
		MPI_Bcast(&width, 1, MPI_INT, 0, MPI_COMM_WORLD);

		MPI_Request req;
		/* Linear segment size */
		sCount = (segSize + 4) * width;
		for(i = 1; i < size - 1; i++) {
			/* Send linear segment size */
			MPI_Send(&sCount, 1, MPI_INT, i, 1, MPI_COMM_WORLD);

			/* Wait for immediate send if this is second iteration */
			if(i > 1) {
				MPI_Wait(&req, &stat);
				free(chunk);
			}

			/* Get subsegment as reference */
			ppm_get_segment_ref(r, rRef, (i - 1) * segSize, 0, segSize + 4, width);
			ppm_get_segment_ref(g, gRef, (i - 1) * segSize, 0, segSize + 4, width);
			ppm_get_segment_ref(b, bRef, (i - 1) * segSize, 0, segSize + 4, width);
			/* Pack everything */
			ppm_pack_n(3, rRef, gRef, bRef, &chunk);

			/* Send pack */
			MPI_Isend(chunk, sCount * 3, MPI_CHAR, i, 1, MPI_COMM_WORLD, &req);

			/* Clear references */
			ppm_clear_ref(rRef);
			ppm_clear_ref(gRef);
			ppm_clear_ref(bRef);
		}
		/* Last node's linear segment size */
		sCount = (lastSegSize + 4) * width;
		/* Send last node's linear segment size */
		MPI_Send(&sCount, 1, MPI_INT, i, 1, MPI_COMM_WORLD);

		/* Wait for immediate send */
		MPI_Wait(&req, &stat);
		free(chunk);

		/* Get subsegment as reference */
		ppm_get_segment_ref(r, rRef, (i - 1) * segSize, 0, lastSegSize + 4, width);
		ppm_get_segment_ref(g, gRef, (i - 1) * segSize, 0, lastSegSize + 4, width);
		ppm_get_segment_ref(b, bRef, (i - 1) * segSize, 0, lastSegSize + 4, width);
		/* Pack everything */
		ppm_pack_n(3, rRef, gRef, bRef, &chunk);

		/* Send pack */
		MPI_Isend(chunk, sCount * 3, MPI_CHAR, i, 1, MPI_COMM_WORLD, &req);

		/* Destroy references and segments */
		ppm_destroy_ref(&rRef);
		ppm_destroy_ref(&gRef);
		ppm_destroy_ref(&bRef);
		ppm_destroy(&r);
		ppm_destroy(&g);
		ppm_destroy(&b);

		/* Wait for immediate send */
		MPI_Wait(&req, &stat);
		free(chunk);

		/* Shrink sizes */
		length -= 4;
		width -= 4;

		/* Write file header */
		FILE *opf = fopen("out.ppm", "w");
		fprintf(opf, "P3\n%d %d\n255\n", width, length);
		fclose(opf);

		byte receivedData;
		char command[30];
		for(i = 1; i < size; i++) {
			/* Receive ACK that process is done and dump file available */
			MPI_Recv(&receivedData, 1, MPI_CHAR, i, tag, MPI_COMM_WORLD, &stat);

			/* Concatenate dump file with output file */
			sprintf(command, "cat out.ppm.%03d >> out.ppm", i);
			system(command);

			/* Remove dump file */
			sprintf(command, "rm out.ppm.%03d", i);
			system(command);
		}
    }
	/* Slave Nodes */
	else {
		unsigned int width, length, sCount;
		ppm_seg_t *r, *g, *b, *or, *og, *ob;
		byte *chunk;
		struct timeval tPInit, tPEnd;

		/* Create segments */
		ppm_create(&r);
		ppm_create(&g);
		ppm_create(&b);
		ppm_create(&or);
		ppm_create(&og);
		ppm_create(&ob);

		/* Receive broadcasted image width and linear segment size */
		MPI_Bcast(&width, 1, MPI_INT, 0, MPI_COMM_WORLD);
		MPI_Recv(&sCount, 1, MPI_INT, 0, 1, MPI_COMM_WORLD, &stat);

		/* Calculate length */
		length = sCount / width;

		/* Receive linear segment */
		chunk = malloc(sCount * 3);
		MPI_Recv(chunk, sCount * 3, MPI_CHAR, 0, 1, MPI_COMM_WORLD, &stat);

		gettimeofday(&tPInit, NULL);

		/* Unpack linear segment into three channels */
		ppm_unpack_n(3, chunk, r, length, width, g, length, width, b, length, width);

		free(chunk);

		/* Apply smooth */
		ppm_smooth(r, or);
		ppm_smooth(g, og);
		ppm_smooth(b, ob);

		/* Remove borders */
		ppm_remove_border(or, PPM_BORDER_TOP, 2);
		ppm_remove_border(or, PPM_BORDER_BOTTOM, 2);
		ppm_remove_border(or, PPM_BORDER_LEFT, 2);
		ppm_remove_border(or, PPM_BORDER_RIGHT, 2);
		ppm_remove_border(og, PPM_BORDER_TOP, 2);
		ppm_remove_border(og, PPM_BORDER_BOTTOM, 2);
		ppm_remove_border(og, PPM_BORDER_LEFT, 2);
		ppm_remove_border(og, PPM_BORDER_RIGHT, 2);
		ppm_remove_border(ob, PPM_BORDER_TOP, 2);
		ppm_remove_border(ob, PPM_BORDER_BOTTOM, 2);
		ppm_remove_border(ob, PPM_BORDER_LEFT, 2);
		ppm_remove_border(ob, PPM_BORDER_RIGHT, 2);

		gettimeofday(&tPEnd, NULL);

		printf("[%d] Elapsed time (only processing): %lu us\n", rank, (((tPEnd.tv_sec * 1000000) + tPEnd.tv_usec) - ((tPInit.tv_sec * 1000000) + tPInit.tv_usec)));

		/* Dump to file */
		char fileName[12];
		sprintf(fileName, "out.ppm.%03d", rank);
		ppm_file_dump(fileName, PPM_FORMAT_P3, or, og, ob);

		/* Send ACK to root node */
		byte done = 1;
		MPI_Send(&done, 1, MPI_CHAR, 0, 1, MPI_COMM_WORLD);

		/* Destroy segments */
		ppm_destroy(&r);
		ppm_destroy(&g);
		ppm_destroy(&b);
		ppm_destroy(&or);
		ppm_destroy(&og);
		ppm_destroy(&ob);
    }

	gettimeofday(&tEnd, NULL);

	if(!rank)
		printf("[%d] Total Elapsed time (including communications): %lu us\n", rank, (((tEnd.tv_sec * 1000000) + tEnd.tv_usec) - ((tInit.tv_sec * 1000000) + tInit.tv_usec)));

	/* Finalise MPI */
    MPI_Finalize();

	return EXIT_SUCCESS;
}